import os
import torch
import csv
from torchvision import transforms
from PIL import Image
import argparse



# Function to load model

def load_model(model_path):
    model = torch.load(model_path)
    model.eval()
    return model

# Function to predict and save results
def predict_and_save(model, data_dir, output_csv):
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
    ])

    with open(output_csv, mode='w', newline='') as file:
        writer = csv.writer(file)
        for img_name in os.listdir(data_dir):
            img_path = os.path.join(data_dir, img_name)
            image = Image.open(img_path).convert('RGB')
            image = transform(image)
            image = image.unsqueeze(0)  # Add batch dimension

            with torch.no_grad():
                outputs = model(image)
                _, predicted = torch.max(outputs, 1)
                label = f"{predicted.item():04d}"

            writer.writerow([img_name, label])

# Main function to parse arguments and run prediction
def main(args):
    data_dir = 'd:/competition/2026quanqiu/data/webfg400_test_A' if args.data_type == 1 else 'd:/competition/2026quanqiu/data/webinat5000_test_A'
    output_csv = 'pred_results_web400.csv' if args.data_type == 1 else 'pred_results_web5000.csv'

    model = load_model(args.model_path)
    predict_and_save(model, data_dir, output_csv)

if __name__ == "__main__":
    # Remove main function
    pass
